# utils.py

import os
import json
import time
import logging
import requests
import numpy as np
from typing import List, Dict, Any, Tuple
from datetime import datetime
import pandas as pd
from sklearn.metrics import cohen_kappa_score, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import config


def setup_logging(log_dir: str) -> logging.Logger:
        """设置日志记录"""
        logger = logging.getLogger('aes_system')
        logger.setLevel(logging.INFO)

        # 创建日志目录
        os.makedirs(log_dir, exist_ok = True)

        # 文件处理器
        log_file = os.path.join(log_dir, f'aes_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log')
        file_handler = logging.FileHandler(log_file, encoding = 'utf-8')
        file_handler.setLevel(logging.INFO)

        # 控制台处理器
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)

        # 设置格式
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        file_handler.setFormatter(formatter)
        console_handler.setFormatter(formatter)

        logger.addHandler(file_handler)
        logger.addHandler(console_handler)

        return logger


def chat(model: str, prompt: str, temperature: float = 0.1, n: int = 1,
         top_p: float = 1, max_tokens: int = -1, timeout: int = 60) -> List[str]:
        """调用LLM API"""
        messages = [{"role": "user", "content": prompt}]
        payload = {
                "messages": messages,
                "model": "glm-4-air",  # 目前使用GLM-4
                "temperature": temperature,
                "n": n,
                "top_p": top_p,
                "max_tokens": max_tokens
        }

        retries = 0
        while retries < 3:  # 最多重试3次
                try:
                        response = requests.post(
                                'https://open.bigmodel.cn/api/paas/v4/chat/completions',
                                headers = {
                                        "Authorization": f"Bearer {config.zhipu_key}",
                                        "Content-Type": "application/json"
                                },
                                json = payload,
                                timeout = timeout
                        )

                        if response.status_code == 200:
                                result = response.json()
                                # 记录token使用情况
                                with open('token_usage.log', 'a', encoding = 'utf-8') as f:
                                        f.write(f"{datetime.now()}: {result['usage']['total_tokens']}\n")
                                return [choice['message']['content'] for choice in result['choices']]

                        retries += 1
                        time.sleep(1)

                except requests.exceptions.RequestException as e:
                        logging.error(f"API request failed: {e}")
                        retries += 1
                        time.sleep(1)

        raise Exception("Failed to get response from LLM API")


def analyze_scoring_results(human_scores: List[int], llm_scores: List[int],
                            save_dir: str) -> Dict[str, float]:
        """分析评分结果"""
        # 基本指标
        metrics = {
                'kappa': cohen_kappa_score(human_scores, llm_scores, weights = 'quadratic'),
                'exact_match': np.mean(np.array(human_scores) == np.array(llm_scores)),
                'within_one': np.mean(abs(np.array(human_scores) - np.array(llm_scores)) <= 1),
                'mean_diff': np.mean(abs(np.array(human_scores) - np.array(llm_scores))),
                'std_diff': np.std(abs(np.array(human_scores) - np.array(llm_scores)))
        }

        # 生成混淆矩阵
        cm = confusion_matrix(human_scores, llm_scores)
        plt.figure(figsize = (10, 8))
        sns.heatmap(cm, annot = True, fmt = 'd', cmap = 'Blues')
        plt.title('Confusion Matrix of Scores')
        plt.xlabel('LLM Scores')
        plt.ylabel('Human Scores')
        plt.savefig(os.path.join(save_dir, 'confusion_matrix.png'))
        plt.close()

        # 生成分数分布对比
        plt.figure(figsize = (10, 6))
        plt.hist(human_scores, alpha = 0.5, label = 'Human Scores', bins = 20)
        plt.hist(llm_scores, alpha = 0.5, label = 'LLM Scores', bins = 20)
        plt.title('Score Distribution Comparison')
        plt.xlabel('Scores')
        plt.ylabel('Frequency')
        plt.legend()
        plt.savefig(os.path.join(save_dir, 'score_distribution.png'))
        plt.close()

        return metrics


def save_scoring_results(essays: List[str], human_scores: List[int],
                         llm_scores: List[int], explanations: List[str],
                         save_path: str):
        """保存评分结果"""
        results_df = pd.DataFrame({
                'essay': essays,
                'human_score': human_scores,
                'llm_score': llm_scores,
                'explanation': explanations,
                'diff': np.array(llm_scores) - np.array(human_scores)
        })

        results_df.to_excel(save_path, index = False)


def parse_json_response(response: str) -> Dict:
        """解析LLM返回的JSON响应"""
        try:
                # 处理可能的markdown格式
                if response.startswith('```json'):
                        response = response.split('```json')[1]
                if response.endswith('```'):
                        response = response.split('```')[0]

                # 清理和解析JSON
                response = response.strip()
                return json.loads(response)
        except Exception as e:
                logging.error(f"JSON parsing error: {e}\nResponse: {response}")
                return {}


def calculate_evaluation_stats(scores: List[float]) -> Dict[str, float]:
        """计算评估统计信息"""
        return {
                'mean': np.mean(scores),
                'std': np.std(scores),
                'min': np.min(scores),
                'max': np.max(scores),
                'median': np.median(scores)
        }


def plot_evolution_progress(kappa_history: List[float], save_path: str):
        """绘制进化过程的kappa值变化"""
        plt.figure(figsize = (10, 6))
        plt.plot(kappa_history)
        plt.title('Evolution Progress')
        plt.xlabel('Generation')
        plt.ylabel('Kappa Score')
        plt.grid(True)
        plt.savefig(save_path)
        plt.close()